import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import os
import sys
import matplotlib as mpl
from matplotlib import font_manager
from matplotlib.font_manager import FontProperties

# 使用自定义字体
simhei_font_path = '/Users/liuyuzhuo/Library/Fonts/SimHei.ttf'
font = FontProperties(fname=simhei_font_path)
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 设置中文字体
if sys.platform.startswith('win'):  # Windows系统
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置黑体
elif sys.platform.startswith('darwin'):  # macOS系统
    # 尝试常见的中文字体
    for font in ['Heiti TC', 'Apple LiGothic', 'STHeiti', 'Hiragino Sans GB']:
        if font in mpl.rcParams['font.family']:
            plt.rcParams['font.sans-serif'] = [font]
            break
    else:
        # 如果找不到，就使用最通用的英文字体
        plt.rcParams['font.sans-serif'] = ['Arial']
else:  # Linux系统
    plt.rcParams['font.sans-serif'] = ['WenQuanYi Zen Hei', 'Droid Sans Fallback']

# 检查系统是否真的支持中文
try:
    # 尝试渲染中文
    plt.figure(figsize=(1, 1))
    plt.text(0.5, 0.5, '测试')
    plt.close()
    use_chinese = True
    print("支持中文显示")
except:
    use_chinese = False
    print("不支持中文显示，将使用英文")

# 设置样式
sns.set_style('whitegrid')
plt.figure(figsize=(12, 8))

# 加载各模型数据
flower_data = pd.read_csv('logs/flower/epoch_losses.csv')
mixed_data = pd.read_csv('logs/mixed_dataset/epoch_losses.csv')
face_data = pd.read_csv('logs/small_face/epoch_losses.csv')

# 使用中文标签
title = '验证损失对比: 通用模型 vs 专业化模型'
xlabel = '训练轮次'
ylabel = '生成器验证损失'
labels = ['花卉模型 (专业化)', '混合数据集模型 (通用)', '人脸模型 (专业化)']
model_names = ['花卉模型', '混合模型', '人脸模型']
init_text = '初始'
final_text = '最终'
imp_text = '改进率'
ds_text = '数据集规模'
images_text = '张图像'

# 绘制验证损失曲线
plt.plot(flower_data['epoch'], flower_data['g_val_loss'], marker='o', linewidth=2, label=labels[0])
plt.plot(mixed_data['epoch'], mixed_data['g_val_loss'], marker='s', linewidth=2, label=labels[1])
plt.plot(face_data['epoch'], face_data['g_val_loss'], marker='^', linewidth=2, label=labels[2])

# 自定义图表
plt.title(title, fontsize=18, fontproperties=font)
plt.xlabel(xlabel, fontsize=14, fontproperties=font)
plt.ylabel(ylabel, fontsize=14, fontproperties=font)
plt.grid(True, linestyle='--', alpha=0.7)
plt.legend(fontsize=12, prop=font)

# 添加初始和最终损失值的标注
for data, label, color, offset in zip(
    [flower_data, mixed_data, face_data],
    model_names,
    ['tab:blue', 'tab:orange', 'tab:green'],
    [0.3, 0.3, -0.6]
):
    # 初始损失
    plt.annotate(
        f'{init_text}: {data["g_val_loss"][0]:.2f}',
        xy=(1, data['g_val_loss'][0]),
        xytext=(1.2, data['g_val_loss'][0] + offset),
        color=color,
        fontweight='bold',
        fontproperties=font
    )
    
    # 最终损失
    plt.annotate(
        f'{final_text}: {data["g_val_loss"][9]:.2f}',
        xy=(10, data['g_val_loss'][9]),
        xytext=(9.7, data['g_val_loss'][9] + offset),
        color=color,
        fontweight='bold',
        fontproperties=font
    )

# 计算改进率
flower_improvement = (flower_data['g_val_loss'][0] - flower_data['g_val_loss'][9]) / flower_data['g_val_loss'][0] * 100
mixed_improvement = (mixed_data['g_val_loss'][0] - mixed_data['g_val_loss'][9]) / mixed_data['g_val_loss'][0] * 100
face_improvement = (face_data['g_val_loss'][0] - face_data['g_val_loss'][9]) / face_data['g_val_loss'][0] * 100

# 添加改进率文本框
textstr = '\n'.join((
    f'{imp_text}:',
    f'{model_names[0]}: {flower_improvement:.2f}%',
    f'{model_names[1]}: {mixed_improvement:.2f}%',
    f'{model_names[2]}: {face_improvement:.2f}%'
))

# 添加数据集大小信息
datasize_str = '\n'.join((
    f'{ds_text}:',
    f'{model_names[0]}: 6,551 {images_text}',
    f'{model_names[1]}: 8,951 {images_text}', 
    f'{model_names[2]}: 2,400 {images_text}'
))

props = dict(boxstyle='round', facecolor='white', alpha=0.8)
plt.text(1.5, 22, textstr, fontsize=12, verticalalignment='top', bbox=props, fontproperties=font)
plt.text(5.5, 22, datasize_str, fontsize=12, verticalalignment='top', bbox=props, fontproperties=font)

# 保存图表
plt.tight_layout()
plt.savefig('model_comparison_loss_curve_zh.png', dpi=300)
print("图表已保存为 'model_comparison_loss_curve_zh.png'") 